package mq.publisher.confirms;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import mq.Constants.Constants;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.TimeoutException;

/**
 * @author: Themberfue
 * @date: 2025/4/28 15:13
 * @description:
 */
public class PublisherConfirms {
    private static final int MAX_MESSAGE = 10000;

    static Connection createConnection() throws Exception {
        // TODO 1. 建立连接
        ConnectionFactory connectionFactory = new ConnectionFactory();
        connectionFactory.setHost(Constants.HOST);
        connectionFactory.setPort(Constants.PORT);
        connectionFactory.setUsername(Constants.USER_NAME);
        connectionFactory.setPassword(Constants.PASSWORD);
        connectionFactory.setVirtualHost(Constants.VIRTUAL_HOST);
        return connectionFactory.newConnection();
    }

    public static void main(String[] args) throws Exception {
        // * Strategy #1: Publishing Messages Individually
//        publishingMessagesIndividually();

        // * Strategy #2: Publishing Messages in Batches
        publishingMessagesInBatches();

        // * Strategy #3: Handling Publisher Confirms Asynchronously
        handlingPublisherConfirmsAsynchronously();
    }

    private static void handlingPublisherConfirmsAsynchronously() throws Exception {
        try (Connection connection = createConnection()) {
            // TODO 2. 开启信道
            Channel channel = connection.createChannel();

            // TODO 3. 开启发布确认机制
            channel.confirmSelect();

            // TODO 4. 声明队列
            channel.queueDeclare(Constants.P_CONFIRMS_QUEUE3, true, false, false, null);

            SortedSet<Long> sortedSet = Collections.synchronizedSortedSet(new TreeSet<>());
            // TODO 5. 监听来自 Broker 的 ack 或 nack
            channel.addConfirmListener(new ConfirmListener() {
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    if (multiple) {
                        sortedSet.headSet(deliveryTag + 1).clear();
                    } else {
                        sortedSet.remove(deliveryTag);
                    }
                }

                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                    if (multiple) {
                        sortedSet.headSet(deliveryTag + 1).clear();
                    } else {
                        sortedSet.remove(deliveryTag);
                    }
                    // TODO 5.1 根据业务逻辑处理消息重传
                }
            });

            Long start = System.currentTimeMillis();
            // TODO 6. 发送消息
            for (int i = 0; i < MAX_MESSAGE; i++) {
                String msg = "hello Publisher Confirms~：" + i;
                Long ackSeq = channel.getNextPublishSeqNo();
                sortedSet.add(ackSeq);
                channel.basicPublish("", Constants.P_CONFIRMS_QUEUE3, null, msg.getBytes(StandardCharsets.UTF_8));
            }
            while (!sortedSet.isEmpty()) {
//                Thread.sleep(10);
            }
            Long end = System.currentTimeMillis();
            System.out.printf("异步确认策略，消息条数：%d；总耗时：%d ms\n", MAX_MESSAGE, end - start);
        }
    }

    private static void publishingMessagesInBatches() throws Exception {
        try (Connection connection = createConnection()) {
            // TODO 2. 开启信道
            Channel channel = connection.createChannel();

            // TODO 3. 开启发布确认机制
            channel.confirmSelect();

            // TODO 4. 声明队列
            channel.queueDeclare(Constants.P_CONFIRMS_QUEUE2, true, false, false, null);

            Long start = System.currentTimeMillis();
            // TODO 5. 发送消息
            int batchSize = 100, outstandingMessageCnt = 0;
            for (int i = 0; i < MAX_MESSAGE; i++) {
                String msg = "hello Publisher Confirms~：" + i;
                channel.basicPublish("", Constants.P_CONFIRMS_QUEUE2, null, msg.getBytes(StandardCharsets.UTF_8));
                outstandingMessageCnt++;

                if (outstandingMessageCnt >= batchSize) {
                    channel.waitForConfirms(5_000);
                    outstandingMessageCnt = 0;
                }
            }
            if (outstandingMessageCnt > 0) {
                channel.waitForConfirms(5_000);
            }
            Long end = System.currentTimeMillis();
            System.out.printf("批量确认策略，消息条数：%d；总耗时：%d ms\n", MAX_MESSAGE, end - start);
        }
    }

    private static void publishingMessagesIndividually() throws Exception {
        try (Connection connection = createConnection()) {
            // TODO 2. 开启信道
            Channel channel = connection.createChannel();

            // TODO 3. 开启发布确认机制
            channel.confirmSelect();

            // TODO 4. 声明队列
            channel.queueDeclare(Constants.P_CONFIRMS_QUEUE1, true, false, false, null);

            Long start = System.currentTimeMillis();
            // TODO 5. 发布消息
            for (int i = 0; i < MAX_MESSAGE; i++) {
                String msg = "hello Publisher Confirms~：" + i;
                channel.basicPublish("", Constants.P_CONFIRMS_QUEUE1, null, msg.getBytes(StandardCharsets.UTF_8));

                // TODO 5.1 等待 5s 收到来自 broker 的确认消息
                channel.waitForConfirms(5_000);
            }
            Long end = System.currentTimeMillis();
            System.out.printf("单独确认策略，消息条数：%d；总耗时：%d ms\n", MAX_MESSAGE, end - start);
        }
    }
}
